import torch
from torch import nn
from collections.abc import Iterable
from typing import List
from varflow.distributions import Distribution
from varflow.transforms import Transform


class FlowLayer(Distribution):
    def __init__(self):
        super().__init__()
        self.base_dist = None
    
    def assign_base(self, base_dist):
        s = self
        while s.base_dist is not None:
            s = s.base_dist
        s.base_dist = base_dist
        return self

    def __call__(self, *args, mode='assign_base', **kwargs):
        return getattr(self, mode)(*args, **kwargs)


class Flow(FlowLayer):
    """
    Base class for Flow.
    Flows use the forward transforms to transform data to noise.
    The inverse transforms can subsequently be used for sampling.
    These are typically useful as generative models of data.
    """

    def __init__(self, transforms: List[Transform], base_dist: Distribution = None):
        super().__init__()
        self.transform = transforms.pop(0)
        self.base_dist = None
        if transforms:
            self.base_dist = Flow(transforms)
        self(base_dist)

    def log_prob(self, x):
        z, ldj = self.transform(x)
        return self.base_dist.log_prob(z) + ldj

    def sample(self, num_samples):
        x = self.base_dist.sample(num_samples)
        return self.transform.inverse(x)
